"""
通过pytorch实现nms算法
"""
# 题目要求：
# ①　导入相关头文件
import numpy as np
import matplotlib.pyplot as plt


# ②　定义函数，形参包括一个图片的所有预测框，置信度，iou阈值，先去置信度的前200
# ③　保存留下来的box索引
def nms(boxes, thresh=0.7):
    # ④　判断box为空返回空Tensor
    if boxes is None or len(boxes) == 0:
        return []
    # ⑤　设置框的所有坐标
    scores = boxes[:, 4]
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    # ⑥　并行化计算所有框的面积
    area = (x2 - x1 + 1) * (y2 - y1 + 1)

    # ⑦　对置信度升序排列  # Desc
    orders = np.argsort(scores)[::-1]
    uncertain = orders.copy()
    rest = np.array([], dtype=np.int64)
    # ⑧　创建新的框中心坐标、宽和高  # Not needed
    # ⑨　循环进行框的筛选，保存得分最大框索引
    while True:
        if len(uncertain) == 0:
            break
        first = uncertain[0]
        uncertain = uncertain[1:]
        rest = np.append(rest, [first])
        if len(uncertain) == 0:
            break

        # ⑩　存储剩下boxes的信息
        # 11　计算当前最大置信度框和其他框的交集
        # 12　计算IOU值
        max_x1 = np.maximum(x1[first], x1[uncertain])
        min_x2 = np.minimum(x2[first], x2[uncertain])
        max_y1 = np.maximum(y1[first], y1[uncertain])
        min_y2 = np.minimum(y2[first], y2[uncertain])
        xx = np.maximum(min_x2 - max_x1 + 1, 0)
        yy = np.maximum(min_y2 - max_y1 + 1, 0)
        I = xx * yy
        U = area[first] + area[uncertain] - I
        IOU = I / (U + 1e-20)

        # 13　通过IOU阈值筛选框，保留最后的框
        uncertain = uncertain[IOU < thresh]

    return rest


def draw_boxes(boxes):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    plt.plot([x1, x2], [y1, y1], color='k')
    plt.plot([x1, x2], [y2, y2], color='k')
    plt.plot([x1, x1], [y1, y2], color='k')
    plt.plot([x2, x2], [y1, y2], color='k')


if '__main__' == __name__:
    boxes = np.array([
        [100, 100, 210, 210, 0.72],
        [280, 290, 420, 420, 0.8],
        [220, 220, 320, 330, 0.92],
        [105, 90, 220, 210, 0.71],
        [230, 240, 325, 330, 0.81],
        [305, 300, 420, 420, 0.9],
        [215, 225, 305, 328, 0.6],
        [150, 260, 290, 400, 0.99],
        [102, 108, 208, 208, 0.72]])

    rest_idx = nms(boxes)
    rest = boxes[rest_idx]

    spr = 1
    spc = 2
    spn = 0
    plt.figure(figsize=[12, 6])

    SIDE = np.max(boxes[:, :4].ravel()) + 100

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Original')
    draw_boxes(boxes)
    plt.xlim([0, SIDE])
    plt.ylim([SIDE, 0])

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('After NMS')
    draw_boxes(rest)
    plt.xlim([0, SIDE])
    plt.ylim([SIDE, 0])

    plt.show()
